#include "utils.h"

ele_class VectorCommitment::get_g1() {
    return g1_;
}

VectorCommitment::VectorCommitment(int r, nlohmann::json vdsseJson): 
  pairing_(vdsseJson["VC"]["param"].get<std::string>()),
  g1_(pairing_, "G1"),
  g2_(pairing_, "G2"),
  gt_(pairing_, "GT"),
  alpha_(pairing_, "Zr"),
  r_(r) {
    g1_.set_str(vdsseJson["VC"]["g1"].get<std::string>());
    g2_.set_str(vdsseJson["VC"]["g2"].get<std::string>());
    alpha_.set_str(vdsseJson["VC"]["alpha"].get<std::string>());
    a_.push_back(alpha_);
    for (int i = 1; i < r; ++i) {
        a_.push_back(a_.back() * alpha_);
    }
    gt_ = pairing_(g1_ ^ alpha_, g2_ ^ a_.back());
}

ele_class VectorCommitment::updateCommit(ele_class C, int u, ele_class m, ele_class m1) {
    return C * (g1_ ^ ((m1 - m) * a_[u]));
}

std::string H(std::string str) {
    auto& hash = Singleton<SHA256>::Instance();
    return mpz_class(hash(str), 16).get_str();
}

std::string F(std::string key, std::string text) {
    const int LEN = 64;
    auto& hash = Singleton<SHA256>::Instance();
    auto& zuc = Singleton<ZUC>::Instance();
    unsigned char key_hash[32];
    hash(key);
    hash.getHash(key_hash);
    unsigned char text_hash[32];
    hash(text);
    hash.getHash(text_hash);
    zuc.Initialization(key_hash, text_hash);
    char result_bytes[LEN];
    zuc.GenerateKeyStream((unsigned int*)result_bytes, LEN >> 2);
    std::string result_str;
    for (size_t i = 0; i < LEN; i++) {
        static const char dec2hex[17] = "0123456789abcdef";
        result_str += dec2hex[(result_bytes[i] >> 4) & 15];
        result_str += dec2hex[ result_bytes[i]       & 15];
    }
    return mpz_class(result_str, 16).get_str();
}

RsaAcc::RsaAcc(pair_class& pairing): acc(pairing, "Zr"), g(pairing, "Zr") {
    acc.set_str("3");
    g.set_str("3"); // 这里简化生成元的求解
}

ele_class RsaAcc::get_acc() {
    return acc;
}

ele_class RsaAcc::insert(std::string text) {
    mpz_class h_text(H(text));
    mpz_nextprime(h_text.get_mpz_t(), h_text.get_mpz_t());
    ele_class prime_text(this->g.ele);
    prime_text = h_text;
    X.insert(prime_text);

    acc = acc ^ prime_text;
    return acc;
}

ele_class RsaAcc::memWitAdd(std::string text) {
    mpz_class h_text(H(text));
    mpz_nextprime(h_text.get_mpz_t(), h_text.get_mpz_t());
    ele_class prime_text(g);
    prime_text = h_text;
    ele_class res(acc);
    res = 0;
    if (X.find(prime_text) == X.end()) {
        return res;
    }
    res = g;
    for (auto it = X.begin(); it != X.end(); ++it) {
        if (*it != prime_text) {
            res = res ^ (*it);
        }
    }
    return res;
}

bool RsaAcc::checkWit(std::string text, ele_class wit) {
    mpz_class h_text(H(text));
    mpz_nextprime(h_text.get_mpz_t(), h_text.get_mpz_t());
    ele_class prime_text(g);
    prime_text = h_text;
    return acc == (wit ^ prime_text);
}

void Setup(nlohmann::json& vdsseJson, pair_class* &pairing, vector<RsaAcc>& Acc, vector<ele_class>& V, int r, VectorCommitment* &VC) {
    using namespace nlohmann;
    using namespace std;
    json& VC_json = vdsseJson["VC"];
    string param = VC_json["param"].get<string>();
    pairing = new pair_class(param);
    VC = new VectorCommitment(r, vdsseJson);
    vdsseJson["EDB"] = nullptr;
    vdsseJson["UpdateCnt"] = nullptr;
    ele_class key(*pairing, "Zr");
    key.random();
    vdsseJson["k"] = key.to_string();
    vdsseJson["r"] = r;
    vdsseJson["Cv"] = nullptr;

    // vdsseJson["X"].clear();
    for (int i = 0; i < r; ++i) {
        // vdsseJson["X"].push_back(R"([])"_json);
        Acc.push_back(RsaAcc(*pairing));
        V.push_back(Acc[i].acc);
    }
    // vdsseJson["v"].clear();
    // for (int i = 0; i < r; ++i) {
    //     vdsseJson["X"].push_back(3);
    // }
}

void Update(int ind, std::string w, std::string op, nlohmann::json& vdsseJson, pair_class* &pairing, vector<RsaAcc>& Acc, vector<ele_class>& V, VectorCommitment* &VC) {
    using namespace std;
    if (!vdsseJson["UpdateCnt"].contains(w)) {
        vdsseJson["UpdateCnt"][w] = 0;
    }
    int cnt = vdsseJson["UpdateCnt"][w].get<int>();
    ++cnt;
    vdsseJson["UpdateCnt"][w] = cnt;
    int r = vdsseJson["r"].get<int>();
    string k = vdsseJson["k"].get<string>();
    string addr = F(k, w + to_string(cnt) + "0");
    auto& base16 = Singleton<BASE16>::Instance();
    mpz_class val_l(base16.encode(to_string(ind) + op), 16);
    mpz_class val_r(F(k, w + to_string(cnt) + "1"), 16);
    mpz_class value(val_l ^ val_r);
    string value_str = value.get_str();
    int i = mpz_class(mpz_class(addr) % r).get_si();
    // vdsseJson["X"][i].push_back(value_str);
    ele_class acc_i = Acc[i].insert(value_str);
    
    vdsseJson["EDB"][addr] = value_str;
    // vector<ele_class> V;
    // for (auto it = Acc.begin(); it != Acc.end(); ++it) {
    //     V.push_back(it->acc);
    // }
    if (vdsseJson["Cv"] == nullptr) {
        ele_class Cv = VC->commit(V);
        vdsseJson["Cv"] = Cv.to_string();
    } else {
        ele_class C(VC->get_g1());
        C.set_str(vdsseJson["Cv"].get<string>());
        ele_class Cv = VC->updateCommit(C, i, acc_i, V[i]);
        vdsseJson["Cv"] = Cv.to_string();
    }
    V[i] = acc_i;
}

void Search(std::string w, nlohmann::json& vdsseJson, vector<RsaAcc>& Acc, vector<ele_class>& V, VectorCommitment* &VC) {
    using namespace std;
    vector<string> TokenList;
    int r = vdsseJson["r"].get<int>();
    if (!vdsseJson["UpdateCnt"].contains(w)) return;
    int cnt = vdsseJson["UpdateCnt"][w].get<int>();
    for (int i = 1; i <= cnt; ++i) {
        TokenList.push_back(F(vdsseJson["k"].get<string>(), w + to_string(i) + "0"));
    }
    // vector<ele_class> V;
    // for (auto it = Acc.begin(); it != Acc.end(); ++it) {
    //     V.push_back(it->acc);
    // }
    vector<string> ResultList;
    vector<int> u_vec;
    // vector<ele_class> pieList;
    vector<string> witList;
    for (auto it = TokenList.begin(); it != TokenList.end(); ++it) {
        string value;
        if (vdsseJson["EDB"].contains(*it)) {
            value = vdsseJson["EDB"][*it].get<string>();
            ResultList.push_back(value);
        }
        int i = mpz_class(mpz_class(*it) % r).get_si();
        u_vec.push_back(i);
        witList.push_back(Acc[i].memWitAdd(value).to_string());
    }
    ele_class C(VC->get_g1());
    C.set_str(vdsseJson["Cv"].get<string>());
    ele_class pie(VC->aggregate(C, u_vec, V));
    vdsseJson["pie"] = pie.to_string();
    vdsseJson["wit"] = nlohmann::json(witList);
    vdsseJson["result"] = nlohmann::json(ResultList);
    vdsseJson["u"] = nlohmann::json(u_vec);
}

void Verify(std::string w, nlohmann::json& vdsseJson, vector<RsaAcc>& Acc, vector<ele_class>& V, VectorCommitment* &VC) {
    using namespace std;
    using namespace nlohmann;
    ele_class C(VC->get_g1());
    C.set_str(vdsseJson["Cv"].get<string>());
    vector<int> u_vec;
    for (auto it = vdsseJson["u"].begin(); it != vdsseJson["u"].end(); ++it) {
        u_vec.emplace_back(*it);
    }
    // vector<ele_class> V;
    // for (auto it = Acc.begin(); it != Acc.end(); ++it) {
    //     V.push_back(it->acc);
    // }
    ele_class pie(VC->get_g1());
    pie.set_str(vdsseJson["pie"].get<string>());
    
    if (VC->verify(C, u_vec, V, pie)) {
        unordered_set<int> FR;
        json& result = vdsseJson["result"];
        string k = vdsseJson["k"].get<string>();
        int result_size = result.size();
        auto& base16 = Singleton<BASE16>::Instance();
        for (int i = 1; i <= result_size; ++i) {
            mpz_class val_r(F(k, w + to_string(i) + "1"), 16);
            mpz_class res_mpz(result[i].get<string>());
            string ind_op = base16.decode(mpz_class(val_r ^ res_mpz).get_str());
            if (ind_op.back() == 'd') { // add
                FR.insert(atoi(ind_op.substr(0, ind_op.size() - 3).c_str()));
            } else {
                FR.erase(atoi(ind_op.substr(0, ind_op.size() - 3).c_str()));
            }
        }
        // accept
        cout << "accept" << endl;
        // FR
    } else {
        // reject
        cout << "reject" << endl;
    }
}
